package com.heiye.blog.ai.advisor.impl;

import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Role;
import com.heiye.blog.ai.advisor.DashScopeAdvisor;
import com.heiye.blog.ai.domain.dos.ChatMessageDO;
import com.heiye.blog.ai.domain.mapper.ChatMessageMapper;
import com.heiye.blog.ai.model.dto.AIChatRequest;
import com.heiye.blog.ai.model.vo.AiChatReqVO;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.transaction.support.TransactionTemplate;
import reactor.core.publisher.Flux;

import java.time.LocalDateTime;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

/**
 * @author: heiye
 * @date: 2025/11/09 下午11:38
 * @version: v1.0.0
 * @description: 自定义百炼模型打印流式日志 Advisor
 */
@Slf4j
public class CustomDashScopeStreamLoggerAndMessage2DBAdvisor implements DashScopeAdvisor {

    private final ChatMessageMapper chatMessageMapper;
    private final AIChatRequest aiChatRequest;
    private final TransactionTemplate transactionTemplate;

    public CustomDashScopeStreamLoggerAndMessage2DBAdvisor(ChatMessageMapper chatMessageMapper,
                                                           AIChatRequest aiChatRequest,
                                                           TransactionTemplate transactionTemplate) {
        this.chatMessageMapper = chatMessageMapper;
        this.aiChatRequest = aiChatRequest;
        this.transactionTemplate = transactionTemplate;
    }

    /**
     * 处理请求参数
     *
     * @param generationParam
     * @return
     */
    @Override
    public GenerationParam adviseRequest(GenerationParam generationParam) {
        // 打印流式日志无需处理参数
        return generationParam;
    }

    /**
     * 处理流式响应
     *
     * @param fluxResult
     * @return
     */
    @Override
    public Flux<GenerationResult> adviseStream(Flux<GenerationResult> fluxResult) {
        // 对话 UUID
        String chatUuid = aiChatRequest.getChatId();
        // 用户消息
        String userMessage = aiChatRequest.getUserMessage();
        // 是否需要深度思考
        Boolean enableThink = aiChatRequest.getThinkEnable();

        // 创建 AI 流式回答聚合容器（线程安全）
        AtomicReference<StringBuilder> fullContent = new AtomicReference<>(new StringBuilder());

        // 使用原子布尔值跟踪分隔线状态（每个请求独立）
        AtomicBoolean needSeparator = new AtomicBoolean(true);

        // 处理返回后的流
        return fluxResult.doOnNext(response -> {
                    // 获取推理内容
                    String reasoningContent = response.getOutput().getChoices().get(0).getMessage().getReasoningContent();

                    // 推理结束后的正式回答
                    String text = response.getOutput().getChoices().get(0).getMessage().getContent();

                    // 是否是正式回答
                    boolean isTextResponse = false;

                    // 逐块收集内容
                    String chunk = null;

                    if (enableThink && StringUtils.isNotEmpty(reasoningContent)) {
                        chunk = reasoningContent;
                    } else {
                        chunk = text;
                        isTextResponse = true;
                    }

                    log.info("## chunk: {}", chunk);

                    // 在正式回答内容之前，添加一个回车
                    if (enableThink
                            && isTextResponse
                            && needSeparator.compareAndSet(true, false)) {
                        chunk = "\n" + chunk;
                    }

                    // 若 chunk 块不为空，则追加到 fullContent 中
                    if (chunk != null) {
                        fullContent.get().append(chunk);
                    }
                })
                .doOnComplete(() -> {
                    // 流完成后打印完整回答
                    String completeResponse = fullContent.get().toString();
                    log.info("\n==== FULL AI RESPONSE ====\n{}\n========================", completeResponse);

                    // 开启编程式事务
                    transactionTemplate.execute(status -> {
                        try {
                            // 1. 存储用户消息
                            chatMessageMapper.insert(ChatMessageDO.builder()
                                    .chatUuid(chatUuid)
                                    .content(userMessage)
                                    // 用户消息
                                    .role(Role.USER.getValue())
                                    .createTime(LocalDateTime.now())
                                    .build());

                            // 2. 存储 AI 回答
                            chatMessageMapper.insert(ChatMessageDO.builder()
                                    .chatUuid(chatUuid)
                                    .content(completeResponse)
                                    // AI 回答
                                    .role(Role.ASSISTANT.getValue())
                                    .createTime(LocalDateTime.now())
                                    .build());

                            return true;
                        } catch (Exception ex) {
                            // 标记事务为回滚
                            status.setRollbackOnly();
                            log.error("", ex);
                        }

                        return false;
                    });
                })
                .doOnError(error -> {
                    // 出错时打印已收集的部分
                    String partialResponse = fullContent.get().toString();
                    log.error("## Stream 流出现错误，已收集回答如下: {}", partialResponse, error);
                });
    }

    /**
     * 执行顺序
     */
    @Override
    public int getOrder() {
        return 99;
    }
}
